from __future__ import division
import torch
import torch.nn as nn
import torch.nn.init as init
from torch.autograd import Function
from torch.autograd import Variable
from math import sqrt
from itertools import product
import numpy as np


class PriorBox(object):
    def __init__(self, cfg):
        super(PriorBox, self).__init__()
        self.image_size = cfg['min_dim']
        self.variance = cfg['variance'] or [0.1]
        self.feature_maps = cfg['feature_maps']
        self.min_sizes = cfg['min_sizes']
        self.max_sizes = cfg['max_sizes']
        self.steps = cfg['steps']
    
    def forward(self):
        mean = []
        for k, f in enumerate(self.feature_maps):
            x,y = np.meshgrid(np.arange(f),np.arange(f))
            x = x.reshape(-1)
            y = y.reshape(-1)
            for i, j in zip(y,x):
                f_k = self.image_size / self.steps[k]
                # 计算网格的中心
                cx = (j + 0.5) / f_k
                cy = (i + 0.5) / f_k

                # 求短边
                s_k = self.min_sizes[k]/self.image_size
                mean += [cx, cy, s_k, s_k]

                # 求长边
                s_k_prime = self.max_sizes[k]/self.image_size
                mean += [cx, cy, s_k_prime, s_k_prime]

                # 获得长方形
                mean += [cx, cy, s_k_prime*sqrt(2), s_k/sqrt(2)]
                mean += [cx, cy, s_k/sqrt(2), s_k_prime*sqrt(2)]
        output = torch.Tensor(mean).view(-1, 4)
        return output